/* 
 * Copyright 2014 Josef Hertl.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package caetus.pkcs11.key;

import caetus.pkcs11.template.Template;
import java.math.BigInteger;
import java.security.NoSuchAlgorithmException;
import java.security.spec.InvalidKeySpecException;
import java.util.logging.Level;
import java.util.logging.Logger;
import sun.security.pkcs11.wrapper.CK_MECHANISM;
import sun.security.pkcs11.wrapper.PKCS11;
import sun.security.pkcs11.wrapper.PKCS11Constants;
import sun.security.pkcs11.wrapper.PKCS11Exception;

/**
 *
 * @author Josef Hertl
 */
public class PublicKey extends Key {

    private PrivateKey privateKey;

    public PublicKey(long sessionHandle, PKCS11 pkcs11, CipherEnum cipherEnum, PrivateKey privateKey) {
        super(sessionHandle, pkcs11, cipherEnum);
        this.privateKey = privateKey;
        privateKey.setPublicKey(this);
    }

    /**
     * Generates both this public key and the corresponding private key.
     *
     * @param pubTemplate The template for this key.
     * @param privTemplate The template for the corresponding private key.
     */
    public void generateKeyPair(Template pubTemplate, Template privTemplate) throws PKCS11Exception {
        try {
            if (pubTemplate == null) {
                pubTemplate = new Template();
                Template basicPubTemplate = templateMaker.getRSAPublicKeyGenerationTemplate(new BigInteger("257").toByteArray(), 512, true, true);
                pubTemplate.merge(basicPubTemplate);
            }
            if (privTemplate == null) {
                privTemplate = new Template();
                Template basicPrivTemplate = templateMaker.getRSAPrivateKeyGenerationTemplate(true, true, true, true);
                privTemplate.merge(basicPrivTemplate);
            }
            long[] keys = pkcs11.C_GenerateKeyPair(sessionHandle, new CK_MECHANISM(cipherEnum.getKeyGenerationMechanism()), pubTemplate.finishAndGetAttributes(), privTemplate.finishAndGetAttributes());

            keyHandle = keys[0];
            privateKey.setKeyHandleAndInitialize(keys[1]);
            super.initialize();
        } catch (InvalidKeySpecException ex) { //this should never happen
            Logger.getLogger(PublicKey.class.getName()).log(Level.SEVERE, null, ex);
            ex.printStackTrace();
        } catch (NoSuchAlgorithmException ex) {
            Logger.getLogger(PublicKey.class.getName()).log(Level.SEVERE, null, ex);
        }
    }

    @Override
    public byte[] encrypt(byte[] dataToEncrypt) throws PKCS11Exception {
        if (!super.isInitialized()) {
            throw new PKCS11Exception(PKCS11Constants.CKR_GENERAL_ERROR);
//            throw new IllegalStateException("The key is not initialized. Use generate() or create() first.");
        } else {
            byte[] encrypted = new byte[64];
            CK_MECHANISM ck_mechanism = new CK_MECHANISM(cipherEnum.getCipherMechanism());
            pkcs11.C_EncryptInit(sessionHandle, ck_mechanism, keyHandle);
            pkcs11.C_Encrypt(sessionHandle, dataToEncrypt, 0, dataToEncrypt.length, encrypted, 0, encrypted.length);
            return encrypted;
        }
    }

    @Override
    public byte[] decrypt(byte[] dataToDecrypt) throws PKCS11Exception {
        throw new UnsupportedOperationException("Public key cannot be used for decryption!");
    }

    @Override
    public byte[] wrapKey(Key keyToWrap) throws PKCS11Exception {
        if (!super.isInitialized()) {
            throw new PKCS11Exception(PKCS11Constants.CKR_GENERAL_ERROR);
//            throw new IllegalStateException("The key is not initialized. Use generate() or create() first.");
        } else {
            CK_MECHANISM ck_mechanism = new CK_MECHANISM(cipherEnum.getCipherMechanism());
            byte[] wrappedKey = pkcs11.C_WrapKey(sessionHandle, ck_mechanism, keyHandle, keyToWrap.getKeyHandle());
            return wrappedKey;
        }
    }

    @Override
    public SecretKey unwrapSecretKey(byte[] toUnwrap, Template template, CipherEnum cipher) throws PKCS11Exception {
        throw new UnsupportedOperationException("Public key cannot be used for unwrapping!");
    }

    @Override
    public PrivateKey unwrapPrivateKey(byte[] toUnwrap, Template template, PublicKey correspondingPublicKey) throws PKCS11Exception {
        throw new UnsupportedOperationException("Public key cannot be used for unwrapping!");
    }

    /**
     * Creates new key values for this key pair. The template specifications are
     * optional. Templates that are null get replaced with the default
     * templates. Also, if the templates are not complete, they are filled up so
     * that they are complete.
     *
     * @param pubTemplate May be null.
     * @param privTemplate May be null.
     * @throws PKCS11Exception
     */
    public void createKeyPair(Template pubTemplate, Template privTemplate) throws PKCS11Exception {
        try {
            if (pubTemplate == null) {
                Template basicPubTemplate = templateMaker.getBasicRSAPublicKeyTemplate(true, true);
                pubTemplate = new Template();
                pubTemplate.merge(basicPubTemplate);
            }
            if (privTemplate == null) {
                Template basicPrivTemplate = templateMaker.getBasicRSAPrivateKeyTemplate(true, true, true, true);
                privTemplate = new Template();
                privTemplate.merge(basicPrivTemplate);
            }
            keyHandle = pkcs11.C_CreateObject(sessionHandle, pubTemplate.finishAndGetAttributes());
            privateKey.setKeyHandleAndInitialize(pkcs11.C_CreateObject(sessionHandle, privTemplate.finishAndGetAttributes()));
        } catch (InvalidKeySpecException ex) { //this should never happen
            Logger.getLogger(PublicKey.class.getName()).log(Level.SEVERE, null, ex);
            ex.printStackTrace();
        } catch (NoSuchAlgorithmException ex) {
            Logger.getLogger(PublicKey.class.getName()).log(Level.SEVERE, null, ex);
            ex.printStackTrace();
        }
    }

    @Override
    public Key copyObject(Template template) throws PKCS11Exception {
        PublicKey copiedKey = new PublicKey(sessionHandle, pkcs11, cipherEnum, privateKey);
        copiedKey.setKeyHandleAndInitialize(pkcs11.C_CopyObject(sessionHandle, keyHandle, template.finishAndGetAttributes()));
        return copiedKey;
    }

    @Override
    public Key deriveKey(CK_MECHANISM mechanism, Template template) throws PKCS11Exception {
        PublicKey derivedKey = new PublicKey(sessionHandle, pkcs11, cipherEnum, privateKey);//there is no proper key pair, but this should do for the tests
        derivedKey.setKeyHandleAndInitialize(pkcs11.C_DeriveKey(sessionHandle, mechanism, keyHandle, template.finishAndGetAttributes()));
        return derivedKey;
    }

    
    
}
